## write a flask
""" from flask import Flask"""
import tempfile
from flask import  Flask, render_template, request, redirect, url_for, session, flash, jsonify,send_from_directory,Response, stream_with_context,abort, send_file
import os
import time
from pathlib import Path
import pandas as pd
import json
from llm_req import ExcelDataAgent, JsonAgent, Agent
from llm_req import list_agent, load_agent
from loguru import logger
from flask_socketio import SocketIO
import threading
from llm_req import LLM
import requests
import urllib.parse as up
import warnings
from flask import Flask, session
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash
from functools import wraps

warnings.filterwarnings(action="ignore", message=".*ssl.*")
current_path = os.path.dirname(os.path.abspath(__file__))
CACHE = Path("~/.cache/agents").expanduser()
if not CACHE.exists():
    CACHE.mkdir(parents=True, exist_ok=True)

app = Flask(__name__)
socketio = SocketIO(app)
app.secret_key = 'sdfdsfds'
# config static folder in ./dist
# config map static uri /assets to ./dist/assets
app.template_folder = os.path.join(current_path ,'dist' )

# 配置登录cookie
app.config['SESSION_TYPE'] = 'filesystem'  # 使用文件系统存储会话
app.config['SESSION_COOKIE_SECURE'] = True  # 如果使用 HTTPS，设置为 True
app.config['SESSION_COOKIE_HTTPONLY'] = True  # 防止 XSS 攻击

# 配置静态文件路径
app.config['STATIC_FOLDER'] = os.path.join(current_path,'dist','assets')
app.config['STATIC_URL_PATH'] = os.path.join(current_path,'dist','assets')
import pathlib
db_path = str(pathlib.Path("~").expanduser() / ".cache" / "llm_flow_service.db" )
app.config['SQLALCHEMY_DATABASE_URI'] = f'sqlite:///{db_path}'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)
SAFE_CODE = os.urandom(10).hex()
logger.success(f"SAFE CODE:   {SAFE_CODE}")
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(120), nullable=False)
    sessions = db.relationship('Session', backref='user', lazy=True)

class Session(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    session_name = db.Column(db.String(80), nullable=False)
    session_data = db.Column(db.Text, nullable=False)  # 存储 JSON 格式的会话数据


try:
    logger.success("db:"+app.config['SQLALCHEMY_DATABASE_URI'])
    with app.app_context():
        db.create_all()
except Exception as e:
    print("db create:",e)

def get_user_id():
    user_id = request.cookies.get('user_id')
    if not user_id:
        return None
    return user_id




@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    safe_code = data.get('code')
        
    username = data.get('username')
    password = data.get('password')
    
    if not username or not password:
        return jsonify({'error': 'Username and password are required'}), 400
    
    if safe_code != SAFE_CODE:
        if not username or not password:
            return jsonify({'error': 'code error'}), 400
    
    u = User.query.filter_by(username=username).first()
    if User.query.filter_by(username=username).first():
        logger.info(f"user:{u}")
        return jsonify({'error': 'Username already exists'}), 400
    
    new_user = User(username=username, password_hash=generate_password_hash(password))
    db.session.add(new_user)
    db.session.commit()
    
    return jsonify({'message': 'User registered successfully'}), 201

@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    logger.info(f"try to login : {username}:{password}")
    user = User.query.filter_by(username=username).first()
    logger.info(f"found user: {user}")
    if not user or not check_password_hash(user.password_hash, password):
        return jsonify({'error': 'Invalid username or password'}), 401
    
    response = jsonify({'message': 'Login successful', 'user_id': user.id})
    response.set_cookie('user_id', str(user.id), secure=True, httponly=True, samesite='Lax')
    
    
    # 这里可以返回一个会话令牌或直接使用用户 ID
    return response, 200

@app.route('/get_sessions', methods=['GET'])
def get_sessions():
    user_id = get_user_id()
    logger.info(f"get session:{user_id}")
    if not user_id:
        return jsonify({'error': 'User ID is required'}), 400
    
    user = User.query.get(user_id)
    if not user:
        return jsonify({'error': 'User not found'}), 404
    
    sessions = {}
    for session in user.sessions:
        chat = json.loads(session.session_data)
        sessions[session.session_name] = chat
    
    return jsonify({'sessions': sessions}), 200

@app.route('/save_session', methods=['POST'])
def save_session():
    chats = request.get_json()
    user_id = get_user_id()

    if not user_id:
        return jsonify({'error': 'User ID is required'}), 400
    Session.query.filter_by(user_id=user_id).delete()
    for session_id in chats:
    
        session = Session.query.get(session_id)
        session_data = chats[session_id]
        data = json.dumps(session_data)
        if not session:
            session = Session(session_name=session_id, session_data = data, user_id=user_id)
            db.session.add(session)
        else:
            session.session_data = data
    db.session.commit()
        
    return jsonify({'message': 'Session saved successfully'}), 200

@app.route('/load_session', methods=['GET'])
def load_session():
    session_id = request.args.get('session_id')
    
    if not session_id:
        return jsonify({'error': 'Session ID is required'}), 400
    
    session = Session.query.get(session_id)
    if not session:
        return jsonify({'error': 'Session not found'}), 404
    
    return jsonify(json.loads(session.session_data)), 200

# 提供静态文件
@app.route('/assets/<path:filename>')
def serve_static(filename):
    return send_from_directory(app.config['STATIC_FOLDER'], filename)

@app.route('/')
def index():
    return render_template('index.html')


def fetch_event_stream(data):
    try:
        logger.warning(f"start ---------------- :{data}")
        remote_host = data.get('remote')
        id = data.get('id')
        type = data.get('type')
        post_data = data.get('data')
        headers = data.get('headers',{})
        stream = data.get('stream',False)
        cookie = data.get('cookie',None)
        use_post = data.get('usepost',False)
        
        # logger.info(f"redirect remote_host:  id: {id} type: {type}")
        logger.info(f"=> {remote_host} data: {post_data} usepost: {use_post} stream: {stream} headers: {headers} cookie: {cookie}")
        if type == 'llm':
            llm = LLM(remote_host=remote_host)
            for chunk in llm.stream(post_data):
                socketio.emit(id, chunk)
        elif type == "curl":
            logger.info(f" curl: {post_data} => {remote_host}")
            sess = requests.Session()
            if len(headers) > 0:
                sess.headers.update(headers)
            if cookie is not None:
                sess.headers['Cookie'] = cookie
            
            method = sess.post if use_post else sess.get
            if "Content-Type" in headers:
                ct = headers["Content-Type"]
                if ct == "application/json":
                    post_data = json.dumps(post_data)
                elif ct == "multipart/form-data":
                    post_data = requests.MultipartEncoder(post_data)
                else:
                    post_data = up.urlencode(post_data)
            
            with method(remote_host, data=post_data, stream=stream, verify=False) as response:
                if response.status_code == 200:
                    try:
                        for line in response.iter_lines():
                            if line:
                                # 解析每一行数据
                                event_data = line.decode('utf-8')
                                socketio.emit(id, event_data)
                    except Exception as e:
                        logger.error(f'Failed to fetch event stream:{e}')
                else:
                    logger.error(f'Failed to fetch event stream:{response.status_code}')
                    logger.error(f"Failed to fetch event stream: {response.content}")
    except Exception as e:
        logger.info(f'Failed to fetch event stream:{e}')
    finally:
        socketio.emit(id, json.dumps({"end":True}))

@socketio.on('message')
def handle_message(data):
    # 启动一个线程来处理event-stream请求
    logger.info(f"messer ======>>>: {data}")
    try:
        th = threading.Thread(target=fetch_event_stream, args=(data,))
        th.start()
    except Exception as e:
        logger.info(f"Error starting thread: {e}")
    finally:
        logger.info(f"Thread finished")

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    user_id = get_user_id()
    if user_id == None:
        raise Exception(f"Request failed with status code {response.status_code}")
        return
    target_url = 'https://115.236.8.156:6404/v1/chat/completions'
    data = request.get_json()
    headers = dict(request.headers)
    params = request.args.to_dict()
    response = requests.post(target_url, json=data,verify=False,headers=headers, params=params, stream=True)

    @stream_with_context
    def _generater():
        for line in response.iter_lines():
            if line:
                yield line
    
    if response.status_code == 200:
        return Response(_generater(), content_type=response.headers['Content-Type'], status=200)
    else:
        raise Exception(f"Request failed with status code {response.status_code}")


# write a file upload
@app.route('/upload', methods=['POST'])
def upload():
    if request.method == 'POST':
        file = request.files['file']
        # allow only txt md excel files
        if file.filename.endswith('.txt') or file.filename.endswith('.md') or file.filename.endswith('.xlsx')or  file.filename.endswith('.csv'):
            # save file to temp file
            temp_dir = tempfile.gettempdir()
            file.save(os.path.join(temp_dir, file.filename))
            if file.filename.endswith('.xlsx'):
                reply_data = []
                for no,i in enumerate(pd.read_excel(os.path.join(temp_dir, file.filename)).fillna('').iloc):
                    item = i.to_dict()
                    item["id"] = no
                    reply_data.append(item)
                return jsonify({
                    "data":reply_data,
                    "name":file.filename,
                })
            elif file.filename.endswith('.csv'):
                reply_data = [ ]
                for no,i in enumerate(pd.read_csv(os.path.join(temp_dir, file.filename)).fillna('').iloc):
                    item = i.to_dict()
                    item["id"] = no
                    reply_data.append(item)

                # reply json

                return jsonify({
                    "data":reply_data,
                    "name":file.filename,
                })
            else:
                with open(os.path.join(temp_dir, file.filename), 'r') as f:
                    reply_data = f.read()
                    # reply json
                    return jsonify({
                        "data":reply_data,
                        "name":file.filename,
                    })
        else:
            return jsonify({
                "error":"file type not supported"
            })
    else:
        return jsonify({
            "error":"no file uploaded"
        })

@app.route('/list_agent', methods=['POST'])
def post_list_agent():
    if request.method == 'POST':
        # get target, points , example from request in json type
        res = []
        for name in os.listdir(str(CACHE)):
            if name.endswith('.json'):
                with open(str(CACHE / name), 'r') as fp:
                    one = json.load(fp)
                    if one["type"] != "":
                        res.append({
                            "value": one.get("name", name.replace(".json","")),
                            "label": one.get("name", name.replace(".json","")),
                            "type":one["type"],
                        })
        return jsonify(res)
    else:
        return jsonify({
            "error":"no file uploaded"
        })

@app.route('/get_batch', methods=['POST'])
def post_get_batch():
    if request.method == 'POST':
        # get target, points , example from request in json type
        data = request.json
        # logger.info(data)
        target = data.get('target','')
        points = data.get('points',[])
        example = data.get('example','')
        useJson = data.get('useJson', False)
        name = data.get('name','')
        dataitems = data.get("data", [])
        model = data.get('model','')
        model_api = data.get('model_api','')
        batch_size = data.get('batch_size', 10)
        threads = data.get('threads', 1)
        
        if name.strip() == '':
            return jsonify({
                "error":"no name provided"
            })
        if useJson:
            ag = JsonAgent(*points,target=target)
        else:
            ag = Agent(*points,target=target)
        if example.strip() != "":
            ag.set_example(example)
        ag.batch_size = int(batch_size)
        ag.threads = threads
        ag._data = dataitems

        batchs = []
        
        for pre in ag.output(batch_size = ag.batch_size):
            # logger.info(pre)
            batchs.append(pre)
        logger.success(f"data:{len(dataitems)} batch size: {ag.batch_size} data: {len(batchs)}")
        
        # logger.success(ag)
        return jsonify({
            "data":batchs   
        })
    else:
        return jsonify({
            "error":"no file uploaded"
        })
    
@app.route('/get_preview', methods=['POST'])
def post_get_preview():
    if request.method == 'POST':
        # get target, points , example from request in json type
        data = request.json
        logger.info(data)
        target = data.get('target','')
        points = data.get('points',[])
        example = data.get('example','')
        useJson = data.get('useJson', False)
        name = data.get('name','')
        dataitems = data.get("data", [])
        model = data.get('model','')
        model_api = data.get('model_api','')
        
        if name.strip() == '':
            return jsonify({
                "error":"no name provided"
            })
        items = data.get('data',[])
        if useJson:
            ag = JsonAgent(*points,target=target)
        else:
            ag = Agent(*points,target=target)
        # ag.set_target(target)
        if len(items) > 0:
            ag |= items
        # if len(points) > 0:
        #     ag.set_point(*points)
        
        if example.strip() != "":
            ag.set_example(example)
        ag |= dataitems
        logger.success(ag)
        ag.set_llm(model, model_api)
        return jsonify({
            "data": str(ag),
        })
    else:
        return jsonify({
            "error":"not post method"
        })
                # reply json


@app.route('/del_agent', methods=['POST'])
def del_agent():
    if request.method == 'POST':
        # get target, points , example from request in json type
        data = request.json
        logger.info(f"recv:{data}")
        name = data.get('name', '')
        if name != "" and name in list_agent():
            logger.info(f"del agent {name} ....")
            
            path = CACHE/ f"{name}.json"
            if path.exists():
                os.remove(str(path))
                return jsonify({"status":"ok"})
            else:
                return jsonify({
                    "error":"no agent found"
                })
        else:
            return jsonify({
                "error":"no name provided"
            })
        

@app.route('/load_agent', methods=['POST'])
def post_load_agent():
    if request.method == 'POST':
        # get target, points , example from request in json type
        data = request.json
        logger.info(f"recv:{data}")
        name = data.get('name', '')
        if name != "" and name in list_agent():
            logger.info(f"load {name} ....")
            oo = load_agent(name)
            dd = oo.to_dict()
            dd["name"] = name
            if isinstance(oo, JsonAgent):
                dd["useJson"] = True
            logger.info(f"load agent {name} success")
            return jsonify(dd)
        elif name != "":
            path = CACHE/ f"{name}.json"
            if path.exists():
                with open(str(path)) as fp:
                    oo = json.load(fp)
                    return jsonify(oo)
            else:
                return jsonify({
                    "error":"no agent found"
                })
        else:
            return jsonify({
                "error":"no name provided"
            })
        
@app.route('/down_me/<file_name>')
def download_file(file_name):
    file_dir = os.path.join(app.root_path, 'static', 'files')
    file_path = os.path.join(file_dir, file_name)
    
    if not os.path.exists(file_path):
        abort(404)
    
    return send_file(file_path, as_attachment=True)

@app.route("/down_results", methods=['POST'])
def get_down_url_by_xlsx_data():
    if request.method == "POST":
        file_dir = os.path.join(app.root_path, 'static', 'files')
        if not pathlib.Path(file_dir).exists():
            pathlib.Path(file_dir).mkdir(parents=True)
        data = request.json.get("data",[])
        tmpName = str(time.time()) + ".xlsx"
        dd = pd.DataFrame(data)
        print(os.path.join(file_dir,tmpName))
        dd.to_excel(os.path.join(file_dir,tmpName))
        return jsonify({"url":"/down_me/"+tmpName,"name":tmpName})
    else:
        return jsonify({
            "error":"no name provided"
        })

@app.route('/save_agent', methods=['POST'])
def post_make_agent():
    if request.method == 'POST':
        # get target, points , example from request in json type
        data = request.json
        logger.info(data)
        
        target = data.get('target','')
        name = data.get('name','')
        typed = data.get('type','')
        
        if target != "":
            points = data.get('points',[])
            example = data.get('example','')
            useJson = data.get('useJson', False)
            
            model = data.get('model','')
            model_api = data.get('model_api','')
            batch_size = data.get('batch_size', 10)
            threads = data.get('threads', 1)
            
            if name.strip() == '':
                return jsonify({
                    "error":"no name provided"
                })
            items = data.get('data',[])
            if useJson:
                ag = JsonAgent(*points,target=target)
            else:
                ag = Agent(*points,target=target)
            # ag.set_target(target)
            if len(items) > 0:
                ag |= items
            ag.batch_size = batch_size
            ag.threads = threads
            # if len(points) > 0:
            #     ag.set_point(*points)
            if example.strip() != "":
                ag.set_example(example)
            logger.success(ag)
            ag.set_llm(model, model_api)
            ag.save(name)
        
            return jsonify({
                "status":"ok"
            })
        else:
            if name != "":
                path = str(CACHE / (name + ".json"))
                with open(path, 'w') as f:
                    json.dump(data, f)
                    return jsonify({
                        "status":"ok"
                    })
            return jsonify({
                "error":"no name provided"
            })
    else:
        return jsonify({
            "error":"no data provided"
        })





def main():
    import argparse
    parser = argparse.ArgumentParser(description='LLM Flow')
    parser.add_argument('--port', type=int, default=5000, help='port to run the server on')
    args = parser.parse_args()

    if args.port > 0:
        # run threading
        # app.run(debug=True, port=args.port)
        logger.info(f"Starting server on port {args.port}")
        app.run(debug=True, port=args.port,host="0.0.0.0", threaded=True)
    
if __name__ == '__main__':
    main()